from langchain.chains import ConversationChain
from langchain.chains.llm import LLMChain
from langchain.prompts import PromptTemplate
from langchain_community.chat_models import ChatTongyi

import os

os.environ["DASHSCOPE_API_KEY"] = "sk-9d8f1914800e497f8717144e860f99bc"
llm = ChatTongyi()

physics_template = """你是一位非常聪明的物理学教授。 \
您擅长以简洁易懂的方式回答有关物理的问题。 \
当你不知道某个问题的答案时，你就承认你不知道。

这是一个问题:
{input}"""

math_template = """你是一位非常优秀的数学家。 你很擅长回答数学问题。 \
你太棒了，因为你能够将难题分解成各个组成部分，\
回答各个组成部分，然后将它们组合起来回答更广泛的问题。

这是一个问题:
{input}"""

prompt_infos = [
    {
        "name": "physics",
        "description": "适合回答物理问题",
        "prompt_template": physics_template,
    },
    {
        "name": "math",
        "description": "适合回答数学问题",
        "prompt_template": math_template,
    },
]
####################### 链dict对象  #######################
destination_chains = {}
# 遍历提示模板，组装成一个chain，然后根据模板名存储到dict中
for p_info in prompt_infos:
    name = p_info["name"]
    prompt_template = p_info["prompt_template"]
    prompt = PromptTemplate(template=prompt_template, input_variables=["input"])
    chain = LLMChain(llm=llm, prompt=prompt)
    destination_chains[name] = chain

####################### 基础链  #######################
default_chain = ConversationChain(llm=llm, output_key="text")  # 起始链路

from langchain.chains.router.llm_router import LLMRouterChain, RouterOutputParser
from langchain.chains.router.multi_prompt_prompt import MULTI_PROMPT_ROUTER_TEMPLATE

# 提取2个链的参数描述
destinations = [f"{p['name']}: {p['description']}" for p in prompt_infos]  # 列表字典
destinations_str = "\n".join(destinations)  # 字符串

# 构建多模板
router_template = MULTI_PROMPT_ROUTER_TEMPLATE.format(destinations=destinations_str)

router_prompt = PromptTemplate(
    template=router_template,
    input_variables=["input"],
    output_parser=RouterOutputParser(),
)

####################### 路由链  #######################
# 根据每个chain中的模板描述选择合适的chain
router_chain = LLMRouterChain.from_llm(llm, router_prompt)

# 使用
from langchain.chains.router import MultiPromptChain

chain = MultiPromptChain(
    router_chain=router_chain,  # 路由链
    destination_chains=destination_chains,
    default_chain=default_chain,  #
    verbose=True,
)
result = chain.invoke("你好")
print(result)